airbyte/lakes.py (1)

12-18: Consider modernizing the deprecated decorator, wdyt?

The abstractproperty decorator has been deprecated since Python 3.3. You could use the modern approach with @property combined with @abstractmethod instead.

-    @abstractproperty
+    @property
+    @abstractmethod
     def uri_protocol(self) -> str:

examples/run_fast_lake_copy.py (1)

64-64: Consider using a more descriptive bucket name for the demo?

The hardcoded "mybucket" might be confusing in a demo context. Perhaps something like "your-s3-bucket-name" or getting it from environment would be clearer, wdyt?

-    bucket_name="mybucket",
+    bucket_name="your-s3-bucket-name",  # Replace with your actual bucket

airbyte/caches/snowflake.py (1)

71-71: Consider moving the import to a type-checking block?

The linter suggests moving the LakeStorage import into a TYPE_CHECKING block since it's only used for type hints in method signatures.

+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    from airbyte.lakes import LakeStorage
+
-from airbyte.lakes import LakeStorage

airbyte/caches/base.py (1)

23-23: Consider moving the import to a type-checking block?

Similar to the Snowflake file, the linter suggests optimizing the LakeStorage import.

+from typing import TYPE_CHECKING, ...
+
+if TYPE_CHECKING:
+    from airbyte.lakes import LakeStorage
+
-from airbyte.lakes import LakeStorage

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

Learnt from: aaronsteers
PR: airbytehq/PyAirbyte#415
File: examples/run_perf_test_reads.py:117-127
Timestamp: 2024-10-09T19:21:45.994Z
Learning: In `examples/run_perf_test_reads.py`, the code for setting up Snowflake configuration in `get_cache` and `get_destination` cannot be refactored into a shared helper function because there are differences between them.

Learnt from: aaronsteers
PR: airbytehq/PyAirbyte#415
File: examples/run_perf_test_reads.py:117-127
Timestamp: 2024-10-09T19:21:45.994Z
Learning: In `examples/run_perf_test_reads.py`, the code for setting up Snowflake configuration in `get_cache` and `get_destination` cannot be refactored into a shared helper function because there are differences between them.

airbyte/caches/bigquery.py (1)

94-94: Same private member access issue as in SnowflakeCache

Similar to the Snowflake implementation, this accesses _read_processor._fully_qualified(). Should we consider a consistent approach across both cache implementations?

Also applies to: 120-120

airbyte/caches/snowflake.py (2)

110-111: Consider validating AWS credentials before use?

The AWS credentials are retrieved but not validated before being embedded in the SQL statement. If the credentials are None or empty, this could lead to authentication failures that might be harder to debug. Would it help to add a check to ensure the credentials are present before proceeding?

 aws_access_key_id = get_secret("AWS_ACCESS_KEY_ID")
 aws_secret_access_key = get_secret("AWS_SECRET_ACCESS_KEY")
+
+if not aws_access_key_id or not aws_secret_access_key:
+    raise ValueError("AWS credentials are required for S3 lake storage operations")

---

164-165: Consider providing more context in the NotImplementedError message?

The error message could be more helpful by explaining what zero-copy loading means and perhaps suggesting alternatives. What do you think about expanding the message to guide users better?

-raise NotImplementedError("Zero-copy loading is not yet supported in Snowflake.")
+raise NotImplementedError(
+    "Zero-copy loading is not yet supported in Snowflake. "
+    "Please use the standard loading method by setting zero_copy=False."
+)

airbyte/caches/bigquery.py (1)

104-104: Document or implement the zero_copy parameter?

The zero_copy parameter is accepted but not used (marked with ARG002). Should this either be implemented, raise a NotImplementedError like in Snowflake, or be documented as reserved for future use?

airbyte/lakes.py (1)

138-139: Consider using consistent path construction?

The dataset paths use a hardcoded "airbyte_data" prefix, but the class already has methods like get_stream_root_path() for constructing paths. Would it be more consistent to use those methods instead?

-output_path = f"{self.bucket_name}/airbyte_data/{table_name}"
+output_path = f"{self.bucket_name}/{self.get_stream_root_path(table_name)}"

Also applies to: 163-164, 204-205, 225-226

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

Learnt from: aaronsteers
PR: airbytehq/PyAirbyte#415
File: examples/run_perf_test_reads.py:117-127
Timestamp: 2024-10-09T19:21:45.994Z
Learning: In `examples/run_perf_test_reads.py`, the code for setting up Snowflake configuration in `get_cache` and `get_destination` cannot be refactored into a shared helper function because there are differences between them.

Learnt from: aaronsteers
PR: airbytehq/PyAirbyte#415
File: examples/run_perf_test_reads.py:117-127
Timestamp: 2024-10-09T19:21:45.994Z
Learning: In `examples/run_perf_test_reads.py`, the code for setting up Snowflake configuration in `get_cache` and `get_destination` cannot be refactored into a shared helper function because there are differences between them.

Learnt from: aaronsteers
PR: airbytehq/PyAirbyte#281
File: airbyte/caches/bigquery.py:40-43
Timestamp: 2024-10-08T15:34:31.026Z
Learning: The `BigQueryCache.get_arrow_dataset` method should have a docstring that correctly states the reason for the `NotImplementedError` as BigQuery not supporting `to_arrow`, instead of incorrectly mentioning `pd.read_sql_table`.

airbyte/caches/snowflake.py (1)

142-178: Similar concerns here, plus some nice touches!

The same issues from the unload method apply here:

• Private member access with self._read_processor.sql_config.schema_name (line 172)
• Potential SQL injection with unescaped identifiers

However, I like a few things about this implementation:

• Clear handling of the zero_copy parameter with a descriptive NotImplementedError
• CASE_INSENSITIVE column matching is thoughtful for data compatibility
• PURGE = FALSE is a safe default to avoid accidental data loss

Would you consider addressing the private member access and SQL escaping concerns similar to the unload method?

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

Learnt from: aaronsteers
PR: airbytehq/PyAirbyte#415
File: examples/run_perf_test_reads.py:117-127
Timestamp: 2024-10-09T19:21:45.994Z
Learning: In `examples/run_perf_test_reads.py`, the code for setting up Snowflake configuration in `get_cache` and `get_destination` cannot be refactored into a shared helper function because there are differences between them.

Learnt from: aaronsteers
PR: airbytehq/PyAirbyte#415
File: examples/run_perf_test_reads.py:117-127
Timestamp: 2024-10-09T19:21:45.994Z
Learning: In `examples/run_perf_test_reads.py`, the code for setting up Snowflake configuration in `get_cache` and `get_destination` cannot be refactored into a shared helper function because there are differences between them.

Learnt from: aaronsteers
PR: airbytehq/PyAirbyte#411
File: airbyte/cli.py:111-160
Timestamp: 2024-10-06T23:44:31.534Z
Learning: In PyAirbyte, error messages in functions like `_resolve_source_job` in `airbyte/cli.py` are designed to decouple the message text from dynamic values, following a structlog-inspired design. Dynamic values are provided via parameters like `input_value`. This approach helps avoid including PII in the message strings, which may be used in telemetry.

Learnt from: Suraj-Vishwakarma70
PR: airbytehq/PyAirbyte#285
File: airbyte/sources/base.py:0-0
Timestamp: 2024-07-09T06:42:41.304Z
Learning: Ensure consistent naming for attributes in the `Source` class in `airbyte/sources/base.py`, such as renaming `_to_be_selected_stream` to `_to_be_selected_streams`.

examples/run_fast_lake_copy.py (3)

52-52: Consider making the GCP project configurable?

The hardcoded project ID "dataline-integration-testing" might limit reusability of this example. Would it make sense to make this configurable via environment variable while keeping the current value as default for simplicity? wdyt?

-    AIRBYTE_INTERNAL_GCP_PROJECT = "dataline-integration-testing"
+    AIRBYTE_INTERNAL_GCP_PROJECT = ab.get_secret("GCP_PROJECT_ID", default="dataline-integration-testing")

---

111-129: Consider consolidating the duplicate Snowflake configuration?

I notice both cache instances use nearly identical configuration with only the schema name differing. Would a helper function make this cleaner, or do you prefer the explicit approach for demo clarity? wdyt?

+def create_snowflake_cache(snowflake_config: dict, warehouse_name: str, schema_name: str) -> SnowflakeCache:
+    return SnowflakeCache(
+        account=snowflake_config["account"],
+        username=snowflake_config["username"],
+        password=snowflake_config["password"],
+        database=snowflake_config["database"],
+        warehouse=warehouse_name,
+        role=snowflake_config["role"],
+        schema_name=schema_name,
+    )
+
-    snowflake_cache_source = SnowflakeCache(
-        account=snowflake_config["account"],
-        username=snowflake_config["username"],
-        password=snowflake_config["password"],
-        database=snowflake_config["database"],
-        warehouse=warehouse_name,
-        role=snowflake_config["role"],
-        schema_name="fast_lake_copy_source",
-    )
-
-    snowflake_cache_dest = SnowflakeCache(
-        account=snowflake_config["account"],
-        username=snowflake_config["username"],
-        password=snowflake_config["password"],
-        database=snowflake_config["database"],
-        warehouse=warehouse_name,
-        role=snowflake_config["role"],
-        schema_name="fast_lake_copy_dest",
-    )
+    snowflake_cache_source = create_snowflake_cache(snowflake_config, warehouse_name, "fast_lake_copy_source")
+    snowflake_cache_dest = create_snowflake_cache(snowflake_config, warehouse_name, "fast_lake_copy_dest")

---

251-283: Consider adding more specific error context?

The current error handling re-raises the exception which is fine for a demo. Would it be helpful to add a bit more context about which step failed, or is the current approach sufficient for demo purposes? wdyt?

    except Exception as e:
-        print(f"\n❌ Error during execution: {e}")
+        print(f"\n❌ Error during fast lake copy workflow: {e}")
+        print("💡 Check your credentials and ensure all services are accessible.")
        raise

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

Learnt from: aaronsteers
PR: airbytehq/PyAirbyte#415
File: examples/run_perf_test_reads.py:117-127
Timestamp: 2024-10-09T19:21:45.994Z
Learning: In `examples/run_perf_test_reads.py`, the code for setting up Snowflake configuration in `get_cache` and `get_destination` cannot be refactored into a shared helper function because there are differences between them.

Learnt from: aaronsteers
PR: airbytehq/PyAirbyte#415
File: examples/run_perf_test_reads.py:117-127
Timestamp: 2024-10-09T19:21:45.994Z
Learning: In `examples/run_perf_test_reads.py`, the code for setting up Snowflake configuration in `get_cache` and `get_destination` cannot be refactored into a shared helper function because there are differences between them.

Learnt from: aaronsteers
PR: airbytehq/PyAirbyte#411
File: airbyte/cli.py:26-26
Timestamp: 2024-10-11T22:05:15.550Z
Learning: In the PyAirbyte project, when reviewing Python code and encountering a TODO comment without an issue link, I should post a friendly reminder to resolve it before merging, instead of suggesting to add an issue link.

airbyte/caches/snowflake.py (1)

87-88: Fix paired destination name mismatch (BigQuery vs Snowflake).

This cache pairs to DestinationSnowflake but paired_destination_name is set to "destination-bigquery". Should this be "destination-snowflake" to avoid downstream pairing issues, wdyt?

Apply:

-    paired_destination_name: ClassVar[str | None] = "destination-bigquery"
+    paired_destination_name: ClassVar[str | None] = "destination-snowflake"

airbyte/caches/snowflake.py (1)

177-177: Avoid private member access; prefer public config attributes.

Accessing self._read_processor.sql_config.schema_name reaches into private members and risks breaking changes. Could we use the public self.schema_name instead for the default case, wdyt?

-            qualified_table_name = f"{self._read_processor.sql_config.schema_name}.{table_name}"
+            qualified_table_name = f"{self.schema_name}.{table_name}"

Also applies to: 228-228

airbyte/caches/snowflake.py (2)

95-103: Reuse LakeStorage’s prefix utility to reduce duplication.

You’ve re-implemented a per-lake artifact prefix; would it be simpler to just delegate to lake_store.get_artifact_prefix() for consistency with the base abstraction, wdyt?

-    def _get_lake_artifact_prefix(self, lake_store: LakeStorage) -> str:
-        """Get the artifact prefix for this lake storage."""
-        return f"AIRBYTE_LAKE_{lake_store.short_name.upper()}_"
+    def _get_lake_artifact_prefix(self, lake_store: LakeStorage) -> str:
+        """Get the artifact prefix for this lake storage."""
+        return lake_store.get_artifact_prefix()

---

230-235: Remove duplicated assignments (DRY).

qualified_prefix, file_format_name, and stage_name are recomputed here after being set on Lines 217-222. Shall we drop this duplicate block to keep things DRY, wdyt?

-        qualified_prefix = (
-            f"{self.database}.{self.schema_name}" if self.database else self.schema_name
-        )
-        file_format_name = self._get_lake_file_format_name(lake_store)
-        stage_name = self._get_lake_stage_name(lake_store)

airbyte/lakes.py (1)

65-67: Nice utility; could we standardize on this across integrations?

Since get_artifact_prefix() exists here, would you like to refactor Snowflake’s _get_lake_artifact_prefix to call this directly for consistency (and future-proofing), wdyt?

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

Learnt from: aaronsteers
PR: airbytehq/PyAirbyte#415
File: examples/run_perf_test_reads.py:117-127
Timestamp: 2024-10-09T19:21:45.994Z
Learning: In `examples/run_perf_test_reads.py`, the code for setting up Snowflake configuration in `get_cache` and `get_destination` cannot be refactored into a shared helper function because there are differences between them.

Learnt from: Suraj-Vishwakarma70
PR: airbytehq/PyAirbyte#285
File: airbyte/sources/base.py:0-0
Timestamp: 2024-07-09T06:42:41.304Z
Learning: Ensure consistent naming for attributes in the `Source` class in `airbyte/sources/base.py`, such as renaming `_to_be_selected_stream` to `_to_be_selected_streams`.

Learnt from: aaronsteers
PR: airbytehq/PyAirbyte#415
File: examples/run_perf_test_reads.py:117-127
Timestamp: 2024-10-09T19:21:45.994Z
Learning: In `examples/run_perf_test_reads.py`, the code for setting up Snowflake configuration in `get_cache` and `get_destination` cannot be refactored into a shared helper function because there are differences between them.

Learnt from: aaronsteers
PR: airbytehq/PyAirbyte#411
File: airbyte/cli.py:111-160
Timestamp: 2024-10-06T23:44:31.534Z
Learning: In PyAirbyte, error messages in functions like `_resolve_source_job` in `airbyte/cli.py` are designed to decouple the message text from dynamic values, following a structlog-inspired design. Dynamic values are provided via parameters like `input_value`. This approach helps avoid including PII in the message strings, which may be used in telemetry.

examples/run_fast_lake_copy.py (1)

158-166: Fix typo in S3 bucket name (will cause runtime failures).

The bucket string is misspelled as "ab-destiantion-iceberg-us-west-2" in both the log line and the constructor. Assuming the intended name is "ab-destination-iceberg-us-west-2", this will 404. Can we correct the spelling and/or make it configurable, wdyt?

Apply this minimal fix:

-    print("   Using co-located bucket: ab-destiantion-iceberg-us-west-2 (us-west-2)")
+    print("   Using co-located bucket: ab-destination-iceberg-us-west-2 (us-west-2)")

     s3_lake = S3LakeStorage(
-        bucket_name="ab-destiantion-iceberg-us-west-2",
+        bucket_name="ab-destination-iceberg-us-west-2",
         region="us-west-2",
         aws_access_key_id=credentials["aws_access_key_id"],
         aws_secret_access_key=credentials["aws_secret_access_key"],
         short_name="s3_main",  # Custom short name for AIRBYTE_LAKE_S3_MAIN_ artifacts
     )

Optionally, would you like to read the bucket name from an env var (e.g., FAST_LAKE_BUCKET_NAME) to avoid hard-coding and improve portability, as previously suggested? wdyt?

examples/run_fast_lake_copy.py (2)

66-73: Clarify error message for missing GCP credentials.

The exception mentions only DEVIN_GCP_SERVICE_ACCOUNT_JSON, but the code also supports GCP_GSM_CREDENTIALS. Shall we update the message to reference both, wdyt?

Apply this diff:

-    if not gcp_creds:
-        raise ValueError(
-            "DEVIN_GCP_SERVICE_ACCOUNT_JSON environment variable not found"
-        )
+    if not gcp_creds:
+        raise ValueError(
+            "Missing GCP credentials: set DEVIN_GCP_SERVICE_ACCOUNT_JSON or GCP_GSM_CREDENTIALS"
+        )

---

241-249: Remove unused unload_results to keep the example lean.

The collected unload_results are not used. For an examples/* script, simpler is better—can we drop the list and append, wdyt?

Apply this diff:

-    unload_results: list[FastUnloadResult] = []
     for stream_name in streams:
-        unload_results.append(
-            snowflake_cache_source.fast_unload_stream(
-                stream_name=stream_name,
-                lake_store=s3_lake,
-            )
-        )
+        snowflake_cache_source.fast_unload_stream(
+            stream_name=stream_name,
+            lake_store=s3_lake,
+        )

And since FastUnloadResult becomes unused, update the import as well:

-from airbyte.lakes import FastUnloadResult, S3LakeStorage
+from airbyte.lakes import S3LakeStorage

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

examples/run_fast_lake_copy.py (6)

69-72: Clarify error message to mention both supported env vars

You accept either DEVIN_GCP_SERVICE_ACCOUNT_JSON or GCP_GSM_CREDENTIALS, but the error mentions only the former. Shall we make the message explicit? wdyt?

Apply this diff:

-        raise ValueError(
-            "DEVIN_GCP_SERVICE_ACCOUNT_JSON environment variable not found"
-        )
+        raise ValueError(
+            "Neither DEVIN_GCP_SERVICE_ACCOUNT_JSON nor GCP_GSM_CREDENTIALS environment variables were found"
+        )

---

23-23: Avoid hard-failing on Windows: guard resource import

The resource module isn’t available on Windows. Importing it at module import time will crash the script. Shall we guard the import so non-Unix users can still run the demo? wdyt?

Apply this diff:

-import resource
+try:
+    import resource  # type: ignore
+except Exception:
+    resource = None  # type: ignore

---

390-399: Skip RLIMIT tweaks if resource is unavailable

Ties in with the guarded import above. Let’s make the RLIMIT update conditional so the script remains cross-platform. wdyt?

Apply this diff:

-    soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
-    print(f"📁 Current file descriptor limits: soft={soft}, hard={hard}")
-    try:
-        new_soft = min(hard, 65536)
-        resource.setrlimit(resource.RLIMIT_NOFILE, (new_soft, hard))
-        soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
-        print(f"📁 Updated file descriptor limits: soft={soft}, hard={hard}")
-    except (ValueError, OSError) as e:
-        print(f"⚠️  Could not increase file descriptor limit: {e}")
+    if resource is not None:
+        soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
+        print(f"📁 Current file descriptor limits: soft={soft}, hard={hard}")
+        try:
+            new_soft = min(hard, 65536)
+            resource.setrlimit(resource.RLIMIT_NOFILE, (new_soft, hard))
+            soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
+            print(f"📁 Updated file descriptor limits: soft={soft}, hard={hard}")
+        except (ValueError, OSError) as e:
+            print(f"⚠️  Could not increase file descriptor limit: {e}")
+    else:
+        print("📁 Skipping RLIMIT tweaks: resource module unavailable on this platform.")

---

196-201: Drop unused variable assignment from read() call

read_result isn’t used later, and dropping it reduces noise in the example. Shall we simplify? wdyt?

Apply this diff:

-        read_result = source.read(
+        source.read(
             cache=snowflake_cache_source,
             force_full_refresh=True,
             write_strategy="replace",
         )

---

241-249: Remove unused unload_results list

The list is never read. Simplifying the loop makes the example tighter. If you remove it, consider also removing the unused FastUnloadResult import. wdyt?

Apply this diff:

-    unload_results: list[FastUnloadResult] = []
-    for stream_name in streams:
-        unload_results.append(
-            snowflake_cache_source.fast_unload_stream(
-                stream_name=stream_name,
-                lake_store=s3_lake,
-            )
-        )
+    for stream_name in streams:
+        snowflake_cache_source.fast_unload_stream(
+            stream_name=stream_name,
+            lake_store=s3_lake,
+        )

---

5-8: Align docstring with actual example (Snowflake + S3 only)

The docstring mentions BigQuery and GCS, but this script exercises Snowflake with S3. Shall we narrow the wording so expectations match the runnable demo? wdyt?

Apply this diff:

-- Direct bulk operations (Snowflake COPY INTO, BigQuery LOAD DATA FROM)
-- Lake storage as an intermediate layer (S3 and GCS)
+- Direct Snowflake bulk operations (COPY INTO)
+- S3 lake storage as an intermediate layer

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

airbyte/caches/snowflake.py (1)

88-89: Fix incorrect paired_destination_name to match Snowflake.

The cache advertises a BigQuery destination name while using a Snowflake destination config. This is misleading and likely breaks pairing logic.

Proposed fix:

-    paired_destination_name: ClassVar[str | None] = "destination-bigquery"
+    paired_destination_name: ClassVar[str | None] = "destination-snowflake"

Wdyt?

airbyte/caches/snowflake.py (5)

114-118: Prefer explicit type check over hasattr for S3-only path.

Using hasattr(lake_store, "aws_access_key_id") is brittle. Would you switch to an explicit isinstance check against S3LakeStorage for clarity and future-proofing, wdyt?

Sketch (local import to avoid cycles):

     def _setup_lake_artifacts(
         self,
         lake_store: LakeStorage,
     ) -> None:
-        if not hasattr(lake_store, "aws_access_key_id"):
+        # Avoid feature detection by attribute; be explicit about supported lake types.
+        from airbyte.lakes import S3LakeStorage  # local import to avoid cycles
+        if not isinstance(lake_store, S3LakeStorage):
             raise NotImplementedError(
                 "Snowflake lake operations currently only support S3 lake storage"
             )

---

186-192: Normalize lake_path_prefix to avoid accidental double slashes in stage URIs.

If callers pass a leading/trailing slash, the resulting path may contain "//".

Proposed tweak:

-        unload_statement = f"""
+        normalized_prefix = lake_path_prefix.strip("/")
+        unload_statement = f"""
             COPY INTO @{qualified_prefix}.{stage_name}/{lake_path_prefix}/
             FROM {qualified_table_name}
             FILE_FORMAT = {qualified_prefix}.{file_format_name}
             OVERWRITE = TRUE
         """

And swap {lake_path_prefix} for {normalized_prefix} in the COPY path. Wdyt?

---

275-282: Normalize lake_path_prefix before COPY INTO load path.

To avoid // in paths when callers pass leading/trailing slashes, normalize prior to constructing the FROM path.

Proposed tweak:

-        load_statement = f"""
+        normalized_prefix = lake_path_prefix.strip("/")
+        load_statement = f"""
             COPY INTO {qualified_table_name}
-            FROM @{qualified_prefix}.{stage_name}/{lake_path_prefix}/
+            FROM @{qualified_prefix}.{stage_name}/{normalized_prefix}/
             FILE_FORMAT = {qualified_prefix}.{file_format_name}
             MATCH_BY_COLUMN_NAME = CASE_INSENSITIVE
             PURGE = FALSE
         """

Wdyt?

---

125-141: Avoid embedding AWS secrets in SQL; validate/quote identifiers; normalize paths.

Creating stages with inline credentials exposes secrets in query history and logs. Also, unquoted identifiers and unvalidated names risk injection/casing surprises.

Would you consider:

• Supporting Snowflake STORAGE INTEGRATION (preferred) and falling back to CREDENTIALS only if missing?
• Validating identifiers (db/schema/stage/format/table) via a small helper before interpolation?
• Normalizing URIs and path components to avoid accidental double slashes?

Minimal, backwards-compatible change within this block:

-        create_stage_sql = f"""
+        storage_integration = getattr(lake_store, "storage_integration_name", None)
+        creds_clause = (
+            f"STORAGE_INTEGRATION = {storage_integration}"
+            if storage_integration
+            else f"""CREDENTIALS = (
+                AWS_KEY_ID = '{lake_store.aws_access_key_id}'
+                AWS_SECRET_KEY = '{lake_store.aws_secret_access_key}'
+            )"""
+        )
+        create_stage_sql = f"""
             CREATE STAGE IF NOT EXISTS {qualified_prefix}.{stage_name}
             URL = '{lake_store.root_storage_uri}'
-            CREDENTIALS = (
-                AWS_KEY_ID = '{lake_store.aws_access_key_id}'
-                AWS_SECRET_KEY = '{lake_store.aws_secret_access_key}'
-            )
+            {creds_clause}
             FILE_FORMAT = {qualified_prefix}.{file_format_name}
         """

If you’d like, I can also add a simple _safe_ident() helper and wire it through these names, wdyt?

---

171-183: Stop accessing private processor attributes; use public config fields.

Accessing self._read_processor.sql_config.schema_name couples to internals; prefer self.schema_name (already present on the config).

Proposed fix:

-        else:
-            qualified_table_name = f"{self._read_processor.sql_config.schema_name}.{table_name}"
+        else:
+            qualified_table_name = f"{self.schema_name}.{table_name}"

Apply the same change in both fast_unload_table (Line 182) and fast_load_table (Line 265). Wdyt?

Also applies to: 254-266

examples/run_fast_lake_copy.py (1)

158-166: Fix S3 bucket name typo and allow env overrides.

“ab-destiantion…” is misspelled and will fail. For demos, env overrides make this easier to run without code changes. Shall we read S3_BUCKET_NAME and AWS_REGION with safe defaults, wdyt?

Proposed change:

-    print(f"🏞️  [{datetime.now().strftime('%H:%M:%S')}] Setting up S3 lake storage...")
-    print("   Using co-located bucket: ab-destiantion-iceberg-us-west-2 (us-west-2)")
+    print(f"🏞️  [{datetime.now().strftime('%H:%M:%S')}] Setting up S3 lake storage...")
+    bucket_name = os.environ.get("S3_BUCKET_NAME", "ab-destination-iceberg-us-west-2")
+    region = os.environ.get("AWS_REGION", "us-west-2")
+    print(f"   Using co-located bucket: {bucket_name} ({region})")
 
     s3_lake = S3LakeStorage(
-        bucket_name="ab-destiantion-iceberg-us-west-2",
-        region="us-west-2",
+        bucket_name=bucket_name,
+        region=region,
         aws_access_key_id=credentials["aws_access_key_id"],
         aws_secret_access_key=credentials["aws_secret_access_key"],
         short_name="s3_main",  # Custom short name for AIRBYTE_LAKE_S3_MAIN_ artifacts
     )

airbyte/caches/snowflake.py (2)

96-104: Reuse LakeStorage.get_artifact_prefix() instead of duplicating logic.

LakeStorage already provides a canonical get_artifact_prefix(). Can we delegate to it to avoid drift?

Proposed change:

     def _get_lake_artifact_prefix(self, lake_store: LakeStorage) -> str:
-        """Get the artifact prefix for this lake storage."""
-        return f"AIRBYTE_LAKE_{lake_store.short_name.upper()}_"
+        """Get the artifact prefix for this lake storage."""
+        return lake_store.get_artifact_prefix()

This keeps naming centralized and consistent across lake types. Wdyt?

---

267-272: Remove duplicated variables (qualified_prefix, file_format_name, stage_name).

These are recomputed immediately after being computed above; the second block is redundant.

Proposed removal:

-        qualified_prefix = (
-            f"{self.database}.{self.schema_name}" if self.database else self.schema_name
-        )
-        file_format_name = self._get_lake_file_format_name(lake_store)
-        stage_name = self._get_lake_stage_name(lake_store)

Keeps the function concise and avoids confusion, wdyt?

examples/run_fast_lake_copy.py (4)

4-11: Align the docstring with this script’s Snowflake-only flow.

The bullet mentions BigQuery LOAD DATA, but this script demonstrates Snowflake → S3 → Snowflake. For clarity in a demo, shall we drop the BigQuery reference here, wdyt?

Proposed change:

-This script demonstrates 100x performance improvements by using:
-- Direct bulk operations (Snowflake COPY INTO, BigQuery LOAD DATA FROM)
+This script demonstrates 100x performance improvements by using:
+- Direct bulk operations (Snowflake COPY INTO)

---

69-72: Clarify the error message to reflect both accepted env vars.

We check DEVIN_GCP_SERVICE_ACCOUNT_JSON or GCP_GSM_CREDENTIALS. Should the error message mention both to reduce confusion?

Proposed change:

-        raise ValueError(
-            "DEVIN_GCP_SERVICE_ACCOUNT_JSON environment variable not found"
-        )
+        raise ValueError(
+            "Missing Google Cloud credentials. Set DEVIN_GCP_SERVICE_ACCOUNT_JSON or GCP_GSM_CREDENTIALS."
+        )

Wdyt?

---

181-188: Consider deriving expected_record_count from actual records for summary.

If RELOAD_INITIAL_SOURCE_DATA is False or the faker config diverges, NUM_RECORDS might not match actual data, skewing % complete. For a demo, would you set expected_record_count = actual_records after counting in Step 1/skip path, wdyt?

No diff provided to keep the script concise, but happy to propose one if you want it wired through the summary prints.

Also applies to: 361-387

---

440-442: Broaden exception handling for rlimit on unsupported platforms.

resource.setrlimit may raise NotImplementedError on some platforms. Shall we include it to avoid noisy tracebacks in the demo, wdyt?

Proposed change:

-    except (ValueError, OSError) as e:
+    except (ValueError, OSError, NotImplementedError) as e:
         print(f"⚠️  Could not increase file descriptor limit: {e}")

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

examples/run_fast_lake_copy.py (2)

291-295: Missing required lake_path_prefix in fast_unload_stream call (runtime TypeError).

SnowflakeCache.fast_unload_stream requires lake_path_prefix. This will raise a TypeError at runtime. Recommend passing stream_name as the prefix to align with your Step 3 load and printed URIs. wdyt?

         unload_results.append(
             snowflake_cache_source.fast_unload_stream(
                 stream_name=stream_name,
                 lake_store=s3_lake,
+                lake_path_prefix=stream_name,
             )
         )

---

197-206: Fix S3 bucket name typo and allow env overrides for bucket/region.

The bucket name has a typo ("ab-destiantion…"), and making bucket/region configurable eases running the example outside your environment without code changes. Shall we use env vars with sane defaults and fix the print? wdyt?

-    print("   Using co-located bucket: ab-destiantion-iceberg-us-west-2 (us-west-2)")
-
-    s3_lake = CustomS3LakeStorage(
-        path_prefix=unique_path_prefix,
-        bucket_name="ab-destiantion-iceberg-us-west-2",
-        region="us-west-2",
-        aws_access_key_id=credentials["aws_access_key_id"],
-        aws_secret_access_key=credentials["aws_secret_access_key"],
-        short_name="s3_main",  # Custom short name for AIRBYTE_LAKE_S3_MAIN_ artifacts
-    )
+    bucket_name = os.environ.get("S3_BUCKET_NAME", "ab-destination-iceberg-us-west-2")
+    region = os.environ.get("AWS_REGION", "us-west-2")
+    print(f"   Using co-located bucket: {bucket_name} ({region})")
+
+    s3_lake = CustomS3LakeStorage(
+        path_prefix=unique_path_prefix,
+        bucket_name=bucket_name,
+        region=region,
+        aws_access_key_id=credentials["aws_access_key_id"],
+        aws_secret_access_key=credentials["aws_secret_access_key"],
+        short_name="s3_main",  # Custom short name for AIRBYTE_LAKE_S3_MAIN_ artifacts
+    )